from miasm.core.asmblock import AsmCFG,AsmBlock,AsmConstraint

# 记录反混淆行为
class ActionList:
    def __init__(self) -> None:
        # 针对图特征
        self.add_blk=[]
        self.del_blk=[]
        self.add_edge=[]
        self.del_edge=[]
        
        # 特征匹配
        self.match_list=[]
        pass
    def graph_add_blk(self,offset):
        self.add_blk.append(offset)
        return
    def graph_del_blk(self,offset):
        self.del_blk.append(offset)
        return
    def graph_add_edge(self,src,dst):
        self.add_edge.append((src,dst))
        return
    def graph_del_edge(self,src,dst):
        self.del_edge.append((src,dst))
        return
    def rule_add_match(self,offset,rule_name):
        self.match_list((offset,rule_name))
        return
    def str_add_blk(self):
        str_add_blk=[]
        for blk in self.add_blk:
            str_add_blk.append(hex(blk))
        return str_add_blk
    def str_del_blk(self):
        str_del_blk=[]
        for blk in self.del_blk:
            str_del_blk.append(hex(blk))
        return str_del_blk
    def str_add_edge(self):
        str_add_edge=[]
        for edge in self.add_edge:
            str_add_edge.append((hex(edge[0]),hex(edge[1])))
        return str_add_edge
    def str_del_edge(self):
        str_del_edge=[]
        for edge in self.del_edge:
            str_del_edge.append((hex(edge[0]),hex(edge[1])))
        return str_del_edge
    
    def _show(self):
        print(self.add_blk)
        print(self.del_blk)
        print(self.add_edge)
        print(self.del_edge)
        print(self.match_list)
        return

def stop_on_jmp(mdis, cur_bloc, offset_to_dis):
    jmp_instr_check = cur_bloc.lines[-1].name in ['JMP','JZ','JNZ']

    if jmp_instr_check:
        cur_bloc.bto.clear()
        offset_to_dis.clear()
        


# 检查AsmCFG在图特征上的信息
def cfg_features_check(asmcfg:AsmCFG,feat_type,feat):
    res_list=[]
    # 遍历块
    for asmblk in asmcfg.blocks:
        pred_list=asmcfg.predecessors(asmblk.loc_key)
        succ_list=asmcfg.successors(asmblk.loc_key)
        # 判断特征
        if feat_type=='out':
            if len(succ_list)==feat:
                blk_offset=asmcfg.loc_db.get_location_offset(asmblk.loc_key)
                res_list.append(blk_offset)
                print(hex(blk_offset))
            pass
        elif feat_type=='in':
            if len(pred_list)==feat:
                blk_offset=asmcfg.loc_db.get_location_offset(asmblk.loc_key)
                res_list.append(blk_offset)
                print(hex(blk_offset))
            pass
        else:
            pass
    return res_list

# 寻找前驱，直到遇到jmp
# 前一个块为唯一前驱，最后一条指令不是jmp，则前驱为逻辑块
def cfg_search_pred_logic(asmcfg:AsmCFG,offset):
    final_offset=offset
    while True:
        cur_blk=asmcfg.getby_offset(final_offset)
        pred_list=asmcfg.predecessors(cur_blk.loc_key)
        pred_succ_list=asmcfg.successors(pred_list[0])
        if len(pred_list)==1 and len(pred_succ_list)==1:
            pred_blk=asmcfg.loc_key_to_block(pred_list[0])
            last_instr=pred_blk.lines[0]
            
            if last_instr.name =='JMP':
                break
            else:
                final_offset=asmcfg.loc_db.get_location_offset(pred_blk.loc_key)
        else:
            break
    return final_offset

# # 将cfg图中的有唯一前驱和后继的图连接起来
# def cfg_combine_node(asmcfg:AsmCFG):
    